
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm

np.random.seed(42)

N = 500
d = 1
x = np.random.uniform(0, 1, N)

piTrue = np.array([0.7, 0.3])
wTrue = np.array([-2.0, 1.0])
bTrue = np.array([0.5, -0.5])
sigmaTrue = np.array([0.4, 0.3])

z = np.random.choice([0, 1], size=N, p=piTrue)
y = np.array([
    np.random.normal(wTrue[zi] * x[i] + bTrue[zi], sigmaTrue[zi])
    for i, zi in enumerate(z)
])

x = x.reshape(-1, 1)
X_tilde = np.hstack([x, np.ones((N, 1))])

piHat = np.array([0.5, 0.5])
wHat = np.array([1.0, -1.0])
bHat = np.array([0.0, 0.0])
sigmaHat = np.array([np.std(y), np.std(y)])

eps = 1e-4
maxIter = 1000
logLikelihoods = []

for iteration in range(maxIter):
    phi_vals = np.zeros((N, 2))
    gamma = np.zeros((N, 2))

    for k in range(2):
        mu_k = wHat[k] * x.flatten() + bHat[k]
        phi_vals[:, k] = norm.pdf(y, loc=mu_k, scale=sigmaHat[k])
        gamma[:, k] = piHat[k] * phi_vals[:, k]

    gamma /= gamma.sum(axis=1, keepdims=True)

    weighted_sum = np.sum([
        piHat[k] * norm.pdf(y, loc=wHat[k] * x.flatten() + bHat[k], scale=sigmaHat[k])
        for k in range(2)
    ], axis=0)
    log_likelihood = np.sum(np.log(weighted_sum))
    logLikelihoods.append(log_likelihood)

    if iteration > 0 and abs(logLikelihoods[-1] - logLikelihoods[-2]) < eps:
        break

    Nk = gamma.sum(axis=0)
    piHat = Nk / N

    for k in range(2):
        G = np.diag(gamma[:, k])
        A = X_tilde.T @ G @ X_tilde
        B = X_tilde.T @ (gamma[:, k] * y)
        w_aug = np.linalg.solve(A, B)
        wHat[k], bHat[k] = w_aug[0], w_aug[1]

    for k in range(2):
        residuals = y - (wHat[k] * x.flatten() + bHat[k])
        sigmaHat[k] = np.sqrt(np.sum(gamma[:, k] * residuals**2) / Nk[k])

print("Estimated Parameters:")
print("pi:", piHat)
print("w:", wHat)
print("b:", bHat)
print("sigma:", sigmaHat)
print("Iterations:", iteration + 1)

plt.figure()
plt.plot(logLikelihoods, marker='o')
plt.title("Log-Likelihood vs Iteration")
plt.xlabel("Iteration")
plt.ylabel("Log-Likelihood")
plt.grid(True)
plt.tight_layout()
plt.savefig("log_likelihood_plot.png")

x_grid = np.linspace(0, 1, 100)
y_line_1 = wHat[0] * x_grid + bHat[0]
y_line_2 = wHat[1] * x_grid + bHat[1]

plt.figure()
plt.scatter(x, y, alpha=0.4, label="Data")
plt.plot(x_grid, y_line_1, label="Estimated Line 1", color='red')
plt.plot(x_grid, y_line_2, label="Estimated Line 2", color='blue')
plt.title("Mixture of Linear Regressions Fit")
plt.xlabel("x")
plt.ylabel("y")
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig("regression_fit_plot.png")
